import calfem.geometry as cfg
import calfem.mesh as cfm
import calfem.vis as cfv
import numpy as np
from toolkits import gauss_line_int, gauss_tri_int


def main():
    # Number of holes
    nH = 1
    # Mesh size factor
    factor = 20

    # externel boundary ID
    dir_id_1 = 100
    dir_id_2 = 200
    dir_id_3 = 300
    neu_id_1 = 400
    # hole boundaryd ID
    hole_neu_id_1 = 500
    hole_neu_id_2 = 600
    hole_neu_id_3 = 700
    hole_neu_id_4 = 800

    # Geometry definition
    g = cfg.Geometry()
    g.point([0.0, 0.0], ID=0)
    g.point([1.0, 0.0], ID=1)
    g.point([1.0, 1.0], ID=2)
    g.point([0.0, 1.0], ID=3)
    g.point([0.4, 0.4], ID=4)
    g.point([0.6, 0.4], ID=5)
    g.point([0.6, 0.6], ID=6)
    g.point([0.4, 0.6], ID=7)
    g.spline([0, 1], el_on_curve=factor, marker=dir_id_1, ID=0)
    g.spline([1, 2], el_on_curve=factor, marker=neu_id_1, ID=1) 
    g.spline([2, 3], el_on_curve=factor, marker=dir_id_2, ID=2) 
    g.spline([3, 0], el_on_curve=factor, marker=dir_id_3, ID=3)
    g.spline([4, 5], el_on_curve=factor/4, marker=hole_neu_id_1, ID=4)
    g.spline([5, 6], el_on_curve=factor/4, marker=hole_neu_id_2, ID=5)
    g.spline([6, 7], el_on_curve=factor/4, marker=hole_neu_id_3, ID=6)
    g.spline([7, 4], el_on_curve=factor/4, marker=hole_neu_id_4, ID=7)
    g.surface([0, 1, 2, 3],[[4,5,6,7]])


    mesh = cfm.GmshMesh(g,return_boundary_elements=True)
    mesh.elType = 2          # Degrees of freedom per node.
    mesh.dofsPerNode = 1     # Factor that changes element sizes.
    # mesh.elSizeFactor = 0.5 # Element size Factor

    coords, edof, dofs, bdofs, elementmarkers, boundaryElements  = mesh.create()

    #############################
    #                           #
    # Elements, Edges and Nodes #
    #                           #
    #############################
    element = edof - 1
    node = coords
    edge = dict()
    for i, e in enumerate(element):
        edge[(e[0], e[1])] = i
        edge[(e[1], e[2])] = i
        edge[(e[2], e[0])] = i

    ######################
    #                    #
    #   Externel edges   #
    #                    #
    ######################
    dir_edge_1 = np.array([edge['node-number-list'] for edge in boundaryElements[dir_id_1]])-1
    dir_edge_2 = np.array([edge['node-number-list'] for edge in boundaryElements[dir_id_2]])-1
    dir_edge_3 = np.array([edge['node-number-list'] for edge in boundaryElements[dir_id_3]])-1

    neu_edge_1 = np.array([edge['node-number-list'] for edge in boundaryElements[neu_id_1]])-1

    dir_ids = [dir_id_1, dir_id_2, dir_id_3]
    neu_ids = [neu_id_1]

    dir_edges = set([tuple(np.array(edge['node-number-list'])-1) for marker in dir_ids for edge in boundaryElements[marker]])
    neu_edges = set([tuple(np.array(edge['node-number-list'])-1) for marker in neu_ids for edge in boundaryElements[marker]])

    ######################
    #                    #
    #     Hole edges     #
    #                    #
    ######################
    # first hole
    hole_neu_edge_1 = np.array([edge['node-number-list'] for edge in boundaryElements[hole_neu_id_1]])-1
    hole_neu_edge_2 = np.array([edge['node-number-list'] for edge in boundaryElements[hole_neu_id_2]])-1
    hole_neu_edge_3 = np.array([edge['node-number-list'] for edge in boundaryElements[hole_neu_id_3]])-1
    hole_neu_edge_4 = np.array([edge['node-number-list'] for edge in boundaryElements[hole_neu_id_4]])-1

    hole_neu_edges = [hole_neu_edge_1, hole_neu_edge_2, hole_neu_edge_3, hole_neu_edge_4]

    ccw_hole_neu_edges = np.array([edge[::-1] for edges in hole_neu_edges[::-1] for edge in edges[::-1]])

    ccw_hole_neu_edge_1 = np.array([edge[::-1] for edge in hole_neu_edge_1[::-1]])
    ccw_hole_neu_edge_2 = np.array([edge[::-1] for edge in hole_neu_edge_2[::-1]])
    ccw_hole_neu_edge_3 = np.array([edge[::-1] for edge in hole_neu_edge_3[::-1]])
    ccw_hole_neu_edge_4 = np.array([edge[::-1] for edge in hole_neu_edge_4[::-1]])

    # print(ccw_hole_neu_edges)
    # print(ccw_hole_neu_edge_1)
    # print(ccw_hole_neu_edge_2)
    # print(ccw_hole_neu_edge_3)
    # print(ccw_hole_neu_edge_4)

    #################################
    #                               #
    #  Crossing elements and edges  #
    #                               #
    #################################
    direction = np.array([0, -1])
    Bh = []
    Eh = {}

    Bh.append(tuple(hole_neu_edge_1[0]))

    # crossing flag 
    connected = False

    ######################################################
    # TODO: New structures of Eh, Bh for multiple holes. #
    ######################################################
    while not connected:
        cross_edge = Bh[-1]
        # print(cross_edge)
        # print(node[cross_edge[0]])
        # print(node[cross_edge[1]])
        if (cross_edge not in dir_edges) and (cross_edge not in neu_edges):
            cross_edge = cross_edge[::-1]
            ei = edge[cross_edge]
            
            i, j, k = element[ei]
            order = [i,j,k]  
            
            local_edges = set([(i,j), (j,k), (k,i)])
            local_edges.remove(cross_edge)
            e1 = np.array(local_edges.pop())
            e2 = np.array(local_edges.pop())
            
            e1n = node[e1[1]] - node[e1[0]]
            e1n = e1n/np.linalg.norm(e1n)
            
            e2n = node[e2[1]] - node[e2[0]]
            e2n = e2n/np.linalg.norm(e2n)
            
            d1 = np.abs(np.dot(direction, e1n))
            d2 = np.abs(np.dot(direction, e2n))
            if d1 < d2:
                ae = np.zeros(3)
                L1 = np.linalg.norm(node[cross_edge[0]] - node[cross_edge[1]])
                L2 = np.linalg.norm(node[e1[0]] - node[e1[1]])
                ae[order.index(cross_edge[0])]= 1/L1
                ae[order.index(e1[0])]= -1/L2
                Eh[ei] = ae
                Bh.append(tuple(e1))
            else:
                ae = np.zeros(3)
                L1 = np.linalg.norm(node[cross_edge[0]] - node[cross_edge[1]])
                L2 = np.linalg.norm(node[e2[0]] - node[e2[1]])
                ae[order.index(cross_edge[0])]= 1/L1
                ae[order.index(e2[0])]= -1/L2
                Eh[ei] = ae
                Bh.append(tuple(e2))
            # print('Element={}, ae={}, next <{},{}>'.format(ei, ae, *Bh[-1]))            
            # print(e1,'\t',e1n)
            # print(e2,'\t',e2n)
            # print(ae)
            # print(d1,d2)
            # print("----------------------------------")
            
                
        else:
            if cross_edge in neu_edges:
                print('Neumann ! ')
            else:
                print('Dirichlet !')
            connected = True
        
    # print(Eh)
            
    # np.all()

    ######################
    #                    #
    #  Numbers of sets   #
    #                    #
    ######################
    E = len(element)
    N = len(node)
    num_dof = N + nH

    # # Display mesh
    # cfv.draw_mesh(coords=coords, edof=edof, dofs_per_node=mesh.dofs_per_node,
    #               el_type=mesh.el_type,filled=True,title="Example 2 - Mesh")
    # cfv.showAndWait()

    # Externel boundary condition
    ########## Dirichlet ##########
    def dir1(x, y):
        return np.zeros_like(x)

    def dir2(x, y):
        return np.sin(x)/np.sin(1)

    def dir3(x, y):
        return np.zeros_like(x)
    ########## Neumann ############
    def neu1(x, y):
        return (1/np.tan(1))*np.sinh(y)/np.sinh(1)

    # Hole's boundary condition
    ########## Dirichlet ##########
    # ...
    ########## Neumann ############
    def hole_neu1(x, y):
        return (np.sin(x)/np.sin(1))*(np.cosh(0.4)/np.sinh(1))

    def hole_neu2(x, y):
        return -(np.cos(0.6)/np.sin(1))*(np.sinh(y)/np.sinh(1))

    def hole_neu3(x, y):
        return -(np.sin(x)/np.sin(1))*(np.cosh(0.6)/np.sinh(1))

    def hole_neu4(x, y):
        return (np.cos(0.4)/np.sin(1))*(np.sinh(y)/np.sinh(1))


    F = np.zeros((num_dof, num_dof))

    for i, e in enumerate(element):
        # print(e)
        # print(node.shape)
        x1,y1 = node[e][0]
        x2,y2 = node[e][1]
        x3,y3 = node[e][2]
        
        A_e = 1/2*np.linalg.det(np.c_[node[e], np.ones((3,1))])
        
        Phi_e = (1/(2*A_e))*np.array([[x2-x3,x3-x1,x1-x2],
                                      [y2-y3,y3-y1,y1-y2]])
        
        
        if i in Eh:
            L1 = np.sqrt((x2-x1)**2+(y2-y1)**2)
            L2 = np.sqrt((x3-x2)**2+(y3-y2)**2)

            left = np.linalg.inv([[(y1-y2)/L1,(x2-x1)/L1],
                                  [(y2-y3)/L2,(x3-x2)/L2]])

            Phi_L = np.c_[left, np.zeros((2,1))]
            
            # print(Eh[i])
            Phi_e = np.c_[Phi_e, Phi_L @ Eh[i]]
            F_e = A_e * Phi_e.T @ Phi_e
            
            ### Todo: multiple holes ###
            ey = np.append(e, -1)

            F[np.ix_(ey,ey)] = F[np.ix_(ey,ey)] + F_e
            # print(i, ey)
        else:
            F_e = A_e * Phi_e.T @ Phi_e
            
            F[np.ix_(e,e)] = F[np.ix_(e,e)] + F_e

    FFF = F.copy()
    print('size: {} x {}'.format(*F.shape))
    print('rank: {}'.format(np.linalg.matrix_rank(F)))



    du = np.zeros(num_dof)

    ##########################
    #   Dirichlet boundary   #
    ##########################
    for edge in dir_edge_1:
        i, j = edge
        L = np.linalg.norm(node[i]-node[j])
        up = gauss_line_int(dir1, *node[i], *node[j])
        du[i] = du[i]-up/L
        du[j] = du[j]+up/L

    for edge in dir_edge_2:
        i, j = edge
        L = np.linalg.norm(node[i]-node[j])
        up = gauss_line_int(dir2, *node[i], *node[j])
        du[i] = du[i]-up/L
        du[j] = du[j]+up/L

    for edge in dir_edge_3:
        i, j = edge
        L = np.linalg.norm(node[i]-node[j])
        up = gauss_line_int(dir3, *node[i], *node[j])
        du[i] = du[i]-up/L
        du[j] = du[j]+up/L

    ##########################################
    #    Neumann boundary (without source)   #
    ##########################################
    s_N = np.zeros(num_dof)
    A = np.eye(num_dof)

    num_red = len(neu_edge_1) + len(ccw_hole_neu_edges)
    num_eff = num_dof - num_red

    # 1. Externel edge
    # effective variable (s1) is at first node of first Neumann edge 
    # (edge's nodes are ordered counterclockwise within the local element).
    ext_eff_dofs = [neu_edge_1[0][0]]
    ext_red_dofs = [edge[1] for edge in neu_edge_1] 

    # 2. Hole edge
    # last edge should be excluded for closed Neumann boundary. 
    hole_eff_dofs = [hole_neu_edge_1[0][0]]
    hole_red_dofs = [edge[1] for edge in ccw_hole_neu_edges[:-1]]

    red_dofs = ext_red_dofs + hole_red_dofs
    # print(ext_red_dofs)
    # print(hole_red_dofs)
    print(ext_eff_dofs)
    print(hole_eff_dofs)



    A[np.ix_(ext_red_dofs, ext_eff_dofs)] = 1
    A[np.ix_(hole_red_dofs, hole_eff_dofs)] = 1

    A = np.delete(A, red_dofs, axis=1)

    # Neumann terms
    ####### TODO : crossing edge #########
    # externel (counter-clockwise)
    fl_bar = 0
    for edge in neu_edge_1:
        i, j = edge
        fl_bar = fl_bar + gauss_line_int(neu1, *node[i], *node[j])
        s_N[j] = fl_bar

    # hole (clockwise)
    # last edge should be excluded for closed Neumann boundary. 
    fl_bar = 0
    for edge in ccw_hole_neu_edge_4:
        i, j = edge
        fl_bar = fl_bar + gauss_line_int(hole_neu4, *node[i], *node[j])
        s_N[j] = fl_bar
        
    for edge in ccw_hole_neu_edge_3:
        i, j = edge    
        fl_bar = fl_bar + gauss_line_int(hole_neu3, *node[i], *node[j])
        s_N[j] = fl_bar
        
    for edge in ccw_hole_neu_edge_2:
        i, j = edge
        fl_bar = fl_bar + gauss_line_int(hole_neu2, *node[i], *node[j])
        s_N[j] = fl_bar
        
    for edge in ccw_hole_neu_edge_1[:-1]:
        i, j = edge 
        fl_bar = fl_bar + gauss_line_int(hole_neu1, *node[i], *node[j])
        s_N[j] = fl_bar
        
    # Source 
    # Particular solution q_0
    # def qx_0(x, y):
    #     return x

    # def qy_0(x, y):
    #     return np.zeros_like(x)

    # for i, e in enumerate(element):
    #     x1,y1 = node[e][0]
    #     x2,y2 = node[e][1]
    #     x3,y3 = node[e][2]
        
    #     A_e = 1/2*np.linalg.det(np.c_[node[e], np.ones((3,1))])
        
    #     Phi_e = (1/(2*A_e))*np.array([[x2-x3,x3-x1,x1-x2],
    #                                   [y2-y3,y3-y1,y1-y2]])
        
    #     int_qx = gauss_tri_int(qx_0, x1, y1, x2, y2, x3, y3, degree=2)
    #     int_qy = gauss_tri_int(qy_0, x1, y1, x2, y2, x3, y3, degree=2)
        
    #     par = Phi_e.T @ np.array([int_qx, int_qy])
        
    #     du[e] = du[e] - par

    # Neumann boundary with source.
    # ...


    FU = F.copy()
    dU = du.copy()

    F = A.T @ FU @ A
    d = A.T @ (dU - FU@s_N)

    r1 = 0

    for edge in hole_neu_edge_1:
        i, j = edge
        r1 = r1 + gauss_line_int(hole_neu1, *node[i], *node[j])

    for edge in hole_neu_edge_2:
        i, j = edge
        r1 = r1 + gauss_line_int(hole_neu2, *node[i], *node[j])

    for edge in hole_neu_edge_3:
        i, j = edge
        r1 = r1 + gauss_line_int(hole_neu3, *node[i], *node[j])

    for edge in hole_neu_edge_4:
        i, j = edge
        r1 = r1 + gauss_line_int(hole_neu4, *node[i], *node[j])
        
    print(r1)

    F[-1,:] = 0
    F[-1,-1] = 1
    d[-1] = r1

    F[0,:] = 0
    F[0,0] = 1
    d[0] = 0


    s_eff = np.linalg.inv(F) @ d
    s = A @ s_eff + s_N
    # print(s_eff[-1])
    # print(d[-1])
    # print( np.linalg.norm(F @ s_eff - d))

    q_h = np.zeros((E,2)) 

    for i, e in enumerate(element):
        x1,y1 = node[e][0]
        x2,y2 = node[e][1]
        x3,y3 = node[e][2]
        
        A_e = 1/2*np.linalg.det(np.c_[node[e], np.ones((3,1))])
        
        Phi_e = 1/(2*A_e)*np.array([[x2-x3,x3-x1,x1-x2],
                                    [y2-y3,y3-y1,y1-y2]])
        
        if i in Eh:
            L1 = np.sqrt((x2-x1)**2+(y2-y1)**2)
            L2 = np.sqrt((x3-x2)**2+(y3-y2)**2)

            left = np.linalg.inv([[(y1-y2)/L1,(x2-x1)/L1],
                                  [(y2-y3)/L2,(x3-x2)/L2]])

            Phi_L = np.c_[left, np.zeros((2,1))]
            
            # print(Phi_e @ s[e])
            Phi_e = np.c_[Phi_e, Phi_L @ Eh[i]]
            
            ey = np.append(e, -1)

            q_e = Phi_e @ s[ey]
            # print(Phi_e @ s[ey])

        else:
            q_e = Phi_e @ s[e]
        # q_e = Phi_e @ s[e]    
        q_h[i] = q_e

    try:
        cfv.figure()

        # Draw the mesh.

        cfv.drawElementValues(
            ev=q_h[:,1],
            coords=coords,
            edof=edof,
            dofs_per_node=mesh.dofsPerNode,
            el_type=mesh.elType,
            title="Example 01",
            draw_elements=True
                )
        cfv.showAndWait()
    except Exception as e:
        pass


if __name__ == '__main__':
    main()